import random
import sys
import threading
import time
import zmq


NBR_CLIENTS = 10
NBR_WORKERS = 3


def tprint(msg):
    sys.stdout.write(msg + '\n')
    sys.stdout.flush()


def client_task(name, i):
    ctx = zmq.Context()
    client = ctx.socket(zmq.REQ)
    client.identity = ("%s-Client%s" % (name, i)).encode("utf-8")
    client.connect("ipc://%s-localfe.ipc" % name)
    while True:
        tprint("%s-Client%s: send HELLO" % (name, i))
        client.send(b"HELLO")
        try:
            reply = client.recv()
        except zmq.ZMQError:
            return
        tprint("%s-Client%s: received %s" % (name, i, reply))
        time.sleep(1)


def worker_task(name, i):
    ctx = zmq.Context()
    worker = ctx.socket(zmq.REQ)
    worker.identity = ("%s-Worker%s" % (name, i)).encode("utf-8")
    worker.connect("ipc://%s-localbe.ipc" % name)
    tprint("%s-Worker%s: send READY" % (name, i))
    worker.send(b"READY")
    while True:
        try:
            msg = worker.recv_multipart()
        except zmq.ZMQError:
            return
        tprint("%s-Worker%s: received %s" % (name, i, msg))
        msg[-1] = b"OK"
        tprint("%s-Worker%s: send %s" % (name, i, msg))
        worker.send_multipart(msg)


def main(myself, peers):
    print("I: preparing broker at %s..." % myself)
    ctx = zmq.Context()

    cloudfe = ctx.socket(zmq.ROUTER)
    if not isinstance(myself, bytes):
        ident = myself.encode("utf-8")
    else:
        ident = myself
    cloudfe.identity = ident
    cloudfe.bind("ipc://%s-cloud.ipc" % myself)

    cloudbe = ctx.socket(zmq.ROUTER)
    cloudbe.identity = ident
    for peer in peers:
        tprint("I: connecting to cloud frontend at %s" % peer)
        cloudbe.connect("ipc://%s-cloud.ipc" % peer)

    if not isinstance(peers[0], bytes):
        peers = [peer.encode("utf-8") for peer in peers]

    localfe = ctx.socket(zmq.ROUTER)
    localfe.bind("ipc://%s-localfe.ipc" % myself)
    localbe = ctx.socket(zmq.ROUTER)
    localbe.bind("ipc://%s-localbe.ipc" % myself)

    input("Press Enter when all brokers are started:")

    for i in range(NBR_WORKERS):
        thread = threading.Thread(target=worker_task, args=(myself, i))
        thread.daemon = True
        thread.start()

    for i in range(NBR_CLIENTS):
        thread_c = threading.Thread(target=client_task, args=(myself, i))
        thread_c.daemon = True
        thread_c.start()

    workers = []

    pollerbe = zmq.Poller()
    pollerbe.register(localbe, zmq.POLLIN)
    pollerbe.register(cloudbe, zmq.POLLIN)

    pollerfe = zmq.Poller()
    pollerfe.register(localfe, zmq.POLLIN)
    pollerfe.register(cloudfe, zmq.POLLIN)

    while True:
        try:
            events = dict(pollerbe.poll(1000 if workers else None))
        except zmq.ZMQError:
            break

        msg = None
        if localbe in events:
            msg = localbe.recv_multipart()
            (address, empty), msg = msg[:2], msg[2:]
            workers.append(address)
            print(myself, " localbe recv msg: ", msg, " after workers queue: ", workers)

            if msg[-1] == b"READY":
                msg = None
        elif cloudbe in events:
            msg = cloudbe.recv_multipart()
            print(myself, " cloudbe recv msg: ", msg, "after workers queue: ", workers)
            (address, empty), msg = msg[:2], msg[2:]
        if msg is not None:
            address = msg[0]
            if address in peers:
                print(myself, " cloudfe send msg: ", msg, "after workers queue: ", workers)
                cloudfe.send_multipart(msg)
            else:
                print(myself, " localfe send msg: ", msg, "after workers queue: ", workers)
                localfe.send_multipart(msg)
        while workers:
            events = dict(pollerfe.poll(0))
            reroutable = False
            if cloudfe in events:
                msg = cloudfe.recv_multipart()
                print(myself, " cloudfe recv msg: ", msg)
                reroutable = False
            elif localfe in events:
                msg = localfe.recv_multipart()
                print(myself, " localfe recv msg: ", msg)
                reroutable = True
            else:
                print('pollerfe.poll break')
                break
            print(myself, " reroutable: ", reroutable)
            if reroutable and peers and random.randint(0, 4) == 0:
                msg = [random.choice(peers), b''] + msg
                print(myself, " cloudbe send msg: ", msg, " after workers queue: ", workers)
                cloudbe.send_multipart(msg)
            else:
                msg = [workers.pop(0), b''] + msg
                print(myself, " localbe send msg: ", msg, " after workers queue: ", workers)
                localbe.send_multipart(msg)
        print('workers is empty')


if __name__ == "__main__":
    if len(sys.argv) >= 2:
        main(myself=sys.argv[1], peers=sys.argv[2:])
    else:
        sys.exit(1)